import logging
import pickle
import numpy as np
from nltk.tokenize import TweetTokenizer

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.nn.utils.rnn as rnn_utils

MM_EMBEDDINGS_DIM = 512
MM_HIDDEN_SIZE = 256
MM_MAX_DICT_SIZE = 500

TOKEN_UNK = "#unk"


class Model(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(Model, self).__init__()
        # 特征提取层
        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 64, 5, stride=5),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=2),
            nn.ReLU(),
        )

        # 获取卷积之后的输出size尺寸
        conv_out_size = self._get_conv_out(input_shape)

        # 得到执行动作的策略
        self.policy = nn.Sequential(
            nn.Linear(conv_out_size, n_actions),
        )

        # 得到当前状态下的Q值
        self.value = nn.Sequential(
            nn.Linear(conv_out_size, 1),
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        fx = x.float() / 256
        conv_out = self.conv(fx).view(fx.size()[0], -1)
        return self.policy(conv_out), self.value(conv_out)


class ModelMultimodal(nn.Module):
    def __init__(self, input_shape, n_actions, max_dict_size=MM_MAX_DICT_SIZE):
        super(ModelMultimodal, self).__init__()

        input_shape = (input_shape[2], input_shape[0], input_shape[1])
        # 卷积网路，应该观测的图像环境状态
        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 64, kernel_size=8, stride=2),  # 输出尺寸 64x64
            nn.ReLU(),
            nn.Conv2d(64, 128, kernel_size=4, stride=2),              # 输出尺寸 31x31
            nn.ReLU(),
            nn.Conv2d(128, 256, kernel_size=4, stride=2),             # 输出尺寸 14x14
            nn.ReLU(),
            nn.Conv2d(256, 512, kernel_size=3, stride=2),            # 输出尺寸 7x7
            nn.ReLU(),
            nn.Conv2d(512, 1024, kernel_size=3, stride=2),  # 输出尺寸 4x4
            nn.ReLU(),
        )


        conv_out_size = self._get_conv_out(input_shape)

        # 词嵌入向量
        self.emb = nn.Embedding(max_dict_size, MM_EMBEDDINGS_DIM)
        self.rnn = nn.LSTM(MM_EMBEDDINGS_DIM, MM_HIDDEN_SIZE, batch_first=True)

        # 执行的动作策略
        self.policy = nn.Sequential(
            nn.Linear(conv_out_size + MM_HIDDEN_SIZE*2, n_actions),
        )

        # 对执行动作得到动作q值
        self.value = nn.Sequential(
            nn.Linear(conv_out_size + MM_HIDDEN_SIZE*2, 1),
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def _concat_features(self, img_out, rnn_hidden):
        '''
        img_out：图像状态数据特征采集输入
        rnn_hidden：文本状态数据特征采集输入
        '''
        batch_size = img_out.size()[0]
        if isinstance(rnn_hidden, tuple):
            # todo 什么情况下rnn_hidden是元组
            # 这里是将文本特征的维度转换为（batch_size, -1）的维度
            flat_h = list(map(lambda t: t.view(batch_size, -1), rnn_hidden))
            rnn_h = torch.cat(flat_h, dim=1)
        else:
            rnn_h = rnn_hidden.view(batch_size, -1)
        # todo rnn进行维度转换是一定会成功的吗？原因是什么？
        # 转换为相同的维度后，才可以继续拼接两个特征状态的数据
        return torch.cat((img_out, rnn_h), dim=1)

    def forward(self, x):
        x_img, x_text = x
        # 输入的文本数据是否符合要求的数据类型
        assert isinstance(x_text, rnn_utils.PackedSequence)

        # deal with text data
        # 处理文本，将文本序列传入得到输出和隐藏层状态
        emb_out = self.emb(x_text.data)
        emb_out_seq = rnn_utils.PackedSequence(emb_out, x_text.batch_sizes)
        rnn_out, rnn_h = self.rnn(emb_out_seq)

        # extract image features
        # 观察图像状态
        fx = x_img
        conv_out = self.conv(fx).view(fx.size()[0], -1)

        # 利用组合特征得到预测的动作概率以及动作Q值
        feats = self._concat_features(conv_out, rnn_h)
        return self.policy(feats), self.value(feats)


class MultimodalPreprocessor:
    '''
    能够处理更加复杂观察数据的预处理器
    在本章示例中处理的是图像和文本的数据
    '''
    log = logging.getLogger("MulitmodalPreprocessor")

    def __init__(self, max_dict_size=MM_MAX_DICT_SIZE, device="cpu"):
        '''
        参数的作用 todo

        param max_dict_size:
        param device:
        '''
        self.max_dict_size = max_dict_size
        self.token_to_id = {TOKEN_UNK: 0}
        self.next_id = 1
        self.tokenizer = TweetTokenizer(preserve_case=True)
        self.device = device

    def __len__(self):
        return len(self.token_to_id)

    def __call__(self, batch):
        """
        Convert list of multimodel observations (tuples with image and text string) into the form suitable
        将特定的模型的观察数据分别提取分离出图像和文本信息
        for ModelMultimodal to disgest
        :param batch:

        return: (观测的图像数据，观测的文本描述压缩数据)
        """
        tokens_batch = []
        # 将文本转换为id后，在重新组成观测数据batch
        for img_obs, txt_obs in batch:
            # tokens = self.tokenizer.tokenize(txt_obs)
            tokens = list(txt_obs)
            idx_obs = self.tokens_to_idx(tokens)
            tokens_batch.append((img_obs, idx_obs))
        # sort batch decreasing to seq len
        # 按照文本问题描述长度进行排序
        tokens_batch.sort(key=lambda p: len(p[1]), reverse=True)
        img_batch, seq_batch = zip(*tokens_batch)
        lens = list(map(len, seq_batch))

        # convert data into the target form
        # images
        img_v = torch.FloatTensor(np.array(img_batch)).to(self.device)
        # sequences 将文本序列数据转换压缩为方便训练的形式
        # 压缩数据前应该要构建一个包含所有文本描述的稀疏矩阵
        seq_arr = np.zeros(shape=(len(seq_batch), max(len(seq_batch[0]), 1)), dtype=np.int64)
        for idx, seq in enumerate(seq_batch):
            seq_arr[idx, :len(seq)] = seq
            # Map empty sequences into single #UNK token
            if len(seq) == 0:
                lens[idx] = 1
        seq_v = torch.LongTensor(seq_arr).to(self.device)
        seq_p = rnn_utils.pack_padded_sequence(seq_v, lens, batch_first=True)
        return img_v, seq_p

    def tokens_to_idx(self, tokens):
        '''
        将文本单词token转换为id，如果对应的单词不在字典里面，则使用#UNK表示
        所以这里的转换只是常见单词的转换，应该是为了更加方便的训练才这样做的
        单词太多可能会导致训练太长
        '''
        res = []
        for token in tokens:
            idx = self.token_to_id.get(token)
            if idx is None:
                if self.next_id == self.max_dict_size:
                    self.log.warning("Maximum size of dict reached, token '%s' converted to #UNK token", token)
                    idx = 0
                else:
                    idx = self.next_id
                    self.next_id += 1
                    self.token_to_id[token] = idx
            res.append(idx)
        return res

    def save(self, file_name):
        '''
        保存生成的字典信息
        '''
        with open(file_name, 'wb') as fd:
            pickle.dump(self.token_to_id, fd)
            pickle.dump(self.max_dict_size, fd)
            pickle.dump(self.next_id, fd)

    @classmethod
    def load(cls, file_name):
        '''
        加载生成的字典信息，并创建预处理对象
        '''
        with open(file_name, "rb") as fd:
            token_to_id = pickle.load(fd)
            max_dict_size = pickle.load(fd)
            next_id = pickle.load(fd)

            res = MultimodalPreprocessor(max_dict_size)
            res.token_to_id = token_to_id
            res.next_id = next_id
            return res


def gen_batch(demo_env, batch_size=128):
    batch_obs, batch_act = [], []
    for _ in range(batch_size):
        obs, _ = demo_env.reset()
        batch_obs.append(obs)
        # 使用标准的 gym 动作生成方式
        for i in range(demo_env.max_steps):
            _, action = demo_env.get_target_area()  # 随机生成点击位置
            batch_act.append(action)
            obs, _, _, _, _ = demo_env.step(demo_env.action_space.sample())

            if i < demo_env.max_steps - 1:
                batch_obs.append(obs)

    assert len(batch_obs) == len(batch_act)
    return batch_obs, batch_act


def train_demo(net, optimizer, demo_env, writer, step_idx, preprocessor, device="cpu"):
    """
    Train net on demonstration batch
    通过人工采集的数据进行神经网路的训练

    net: 强化学习网络
    optimier: 优化器
    batch: 人工采集的数据（环境状态，执行的动作）
    writer: tensorboard writer
    step_idx: 训练步数
    preprocessor: 预处理器
    device: 训练所使用的设备是cpu还是gpu
    """
    batch_obs, batch_act = gen_batch(demo_env)
    batch_v = preprocessor(batch_obs)
    optimizer.zero_grad()
    ref_actions_v = torch.LongTensor(batch_act).to(device)
    policy_v = net(batch_v)[0]
    loss_v = F.cross_entropy(policy_v, ref_actions_v)
    loss_v.backward()
    optimizer.step()
    writer.add_scalar("demo_loss", loss_v.item(), step_idx)

